import logging
import torch
from ultralytics import YOLO
import argparse
import os
import cv2
import sys
from datetime import datetime

def check_env(args):
    if args.output_dir is not None:
        try:
            os.makedirs(args.output_dir, exist_ok=True)
            logging.info(f"Create output dir: {args.output_dir}")
        except Exception as e:
            logging.error(f"Create output dir failed: {e}")
            return False
    else:
        logging.error(f"Output dir is None!")
        return False
    
    if torch.cuda.is_available() == False:
        logging.error("CUDA is not available!")
        return False
    
    if os.path.exists(args.model_file) == False:
        logging.error(f"Model path {args.model_file} not exists!")
        return False
    
    return True

def main():

    time_stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    # 日志初始化
    logger = logging.getLogger(__name__)
    logdir = os.path.join(os.path.abspath(os.path.dirname(__file__)), "Logs", __name__, time_stamp)
    os.makedirs(logdir, exist_ok=True)
    log_file = os.path.join(logdir, f"{__name__}.log")
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s",
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler(sys.stdout)
        ]
    )

    try:
        parser = argparse.ArgumentParser(description=__name__)
        parser.add_argument("--output_dir", type=str, default="output", help="Output dir")
        parser.add_argument("--model_file", type=str, default="/home/xlli67/Project/yolo_train/input/model/0816/weights/best.engine", help="Model file")
        parser.add_argument("--input_img", type=str, default="/home/xlli67/Project/yolo_train/input/images/Image_20250327161453.png", help="Input image")
        parser.add_argument("--inference_param_device", type=str, default="cuda", help="Device used in inference, e.g. 'cpu', '0', '0,1,2,3'")
        args = parser.parse_args()
    except Exception as e:
        logging.error(f"Parse args failed: {e}")
        return
    
    if check_env(args) == False:
        logging.error("Check env failed!")
        return
    
    result_save_dir = os.path.join(args.output_dir, "inference", "seg")
    os.makedirs(result_save_dir, exist_ok=True)

    try:
        img = cv2.imread(args.input_img)
        img_tensor = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0)
        logger.info(f"Load image success! Origin image shape: {img.shape}")
    except Exception as e:
        logging.error(f"Load image failed: {e}")
        return
    
    model = YOLO(args.model_file)
    model.fuse()
    model.eval()
    device = None
    try:
        device = torch.device(args.inference_param_device)
    except Exception as e:
        logging.error(f"Device {args.inference_param_device} invalid: {e}")
        return
    model.to(device)
    img_tensor = img_tensor.to(device)

    
    logger.info(f"Start inference with...")
    logger.info(f"input image: {args.input_img}")
    logger.info(f"model file: {args.model_file}")
    logger.info(f"device: {args.inference_param_device}")
    # 开始推理
    results = model.predict(img_tensor,
                            imgsz=960,
                            device=device,
                            retina_masks=True)
    
    for i,result in enumerate(results):
        result.save(os.path.join(result_save_dir, f"result_{i}.jpg"))
        logger.info(f"Save result {i} to {result_save_dir} success!")
